package com.ruoyi.web.core.utils;

import lombok.extern.slf4j.Slf4j;
import lombok.val;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * <pre>
 *
 * </pre>
 *
 * @author zhangli
 * @since 2023/9/24
 */
@Slf4j
public class RsaUtils {

    // 枚举
    public enum KeyType {
        PRIVATE, PUBLIC
    }

    public static final int ENCRYPT_LIMIT = 117;

    /**
     * 把base64格式的公钥转换成java公钥.
     *
     * @param publicKey base64格式的公钥
     * @return PublicKey
     */
    public static PublicKey base64ToPublicKey(String publicKey) {
        return (PublicKey) base64ToKey(publicKey, KeyType.PUBLIC);
    }

    /**
     * 把base64格式的私钥转换成java私钥.
     *
     * @param privateKey base64格式的私钥
     * @return PrivateKey
     */
    public static PrivateKey base64ToPrivateKey(String privateKey) {
        return (PrivateKey) base64ToKey(privateKey, KeyType.PRIVATE);
    }

    /**
     * 把java秘钥转换成base64格式字符串.
     *
     * @param key java秘钥
     * @return String
     */
    public static String keyToBase64(Key key) {
        return Base64.getEncoder().encodeToString(key.getEncoded());
    }

    /**
     * 内容签名.
     *
     * @param content 需要签名的内容
     * @param key     签名私钥
     */
    public static String signature(String content, PrivateKey key) {
        log.info("signature:{} {}", content, key);
        try {
            val signature = Signature.getInstance("SHA1WithRSA");
            signature.initSign(key);
            signature.update(content.getBytes());
            return Base64.getEncoder().encodeToString(signature.sign());
        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException e) {
            throw new SecurityException(e);
        }
    }

    /**
     * 内容验签.
     *
     * @param content 验签内容
     * @param sign    签名
     * @param key     公钥
     */
    public static boolean verify(String content, String sign, String key) {
        PublicKey publicKey = base64ToPublicKey(key);
        return verify(content, sign, publicKey);
    }

    /**
     * 内容验签.
     *
     * @param content 验签内容
     * @param sign    签名
     * @param key     公钥
     */
    public static boolean verify(String content, String sign, PublicKey key) {
        try {
            val signature = Signature.getInstance("SHA1WithRSA");
            signature.initVerify(key);
            signature.update(Base64.getDecoder().decode(content));
            return signature.verify(Base64.getDecoder().decode(sign));
        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException e) {
            throw new SecurityException(e);
        }
    }

    /**
     * 公钥加密.
     *
     * @param content 加密内容
     * @param key     密钥
     */
    public static String encode(String content, Key key) {
        ByteArrayOutputStream outputStream = null;
        try {
            val contentBytes = content.getBytes(StandardCharsets.UTF_8);
            val length = contentBytes.length;
            val cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            if (length <= ENCRYPT_LIMIT) {
                return
                        Base64.getEncoder().encodeToString(cipher.doFinal(contentBytes));
            } else {
                outputStream = new ByteArrayOutputStream();
                for (int i = 0; i < length; i += ENCRYPT_LIMIT) {
                    val size = Math.min((length - i), ENCRYPT_LIMIT);
                    outputStream.write(cipher.doFinal(contentBytes, i, size));
                }
                return Base64.getEncoder().encodeToString(outputStream.toByteArray());
            }
        } catch (IOException
                | NoSuchAlgorithmException
                | InvalidKeyException
                | NoSuchPaddingException
                | BadPaddingException
                | IllegalBlockSizeException e) {
            throw new SecurityException(e);
        } finally {
            if (outputStream != null) {
                try {
                    outputStream.close();
                } catch (IOException ignore) {
                    log.debug(ignore.getMessage(), ignore);
                }
            }
        }
    }

    /**
     * 私钥解密.
     *
     * @param content 加密内容
     * @param key     密钥
     */
    public static String decode(String content, Key key) {
        int maxLimit = ((RSAKey) key).getModulus().bitLength() / 8;
        ByteArrayOutputStream outputStream = null;
        try {
            val cipher = Cipher.getInstance("RSA");
            val contentBytes = Base64.getDecoder().decode(content);
            cipher.init(Cipher.DECRYPT_MODE, key);
            val length = contentBytes.length;
            if (length <= maxLimit) {
                return new String(cipher.doFinal(contentBytes),
                        StandardCharsets.UTF_8);
            } else {
                outputStream = new ByteArrayOutputStream();
                for (int i = 0; i < length; i += maxLimit) {
                    val size = Math.min((length - i), maxLimit);
                    outputStream.write(cipher.doFinal(contentBytes, i, size));
                }
                return new String(outputStream.toByteArray(), StandardCharsets.UTF_8);
            }
        } catch (NoSuchAlgorithmException
                | NoSuchPaddingException
                | IOException
                | InvalidKeyException
                | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new SecurityException(e);
        } finally {
            if (outputStream != null) {
                try {
                    outputStream.close();
                } catch (IOException ignore) {
                    log.debug(ignore.getMessage(), ignore);
                }
            }
        }
    }

    /**
     * 把PEM格式的秘钥格式化成java的.
     *
     * @param key 需要格式化的key 不是pem的也没关系
     */
    public static String pemKeyClear(String key) {
        if (key == null || key.trim().equals("")) {
            throw new IllegalArgumentException(key);
        }
        val keySegments = key.split("\n");
        if (keySegments.length == 1) {
            return keySegments[0];
        } else {
            if (keySegments.length <= 2) {
                throw new IllegalArgumentException("不正确的秘钥格式");
            }
            val javaKey = new StringBuilder();
            for (int i = 1; i < keySegments.length - 1; i++) {
                javaKey.append(keySegments[i].replace("\r", ""));
            }
            return javaKey.toString();
        }
    }

    /**
     * 把BASE64秘钥转换成java秘钥.
     *
     * @param key     需要转换的秘钥
     * @param keyType 秘钥类型
     * @see KeyType
     */
    public static Key base64ToKey(String key, KeyType keyType) {
        assert keyType != null && key != null;
        key = pemKeyClear(key);
        try {
            val keyFactory = KeyFactory.getInstance("RSA");
            if (keyType == KeyType.PUBLIC) {
                return keyFactory
                        .generatePublic(
                                new X509EncodedKeySpec(
                                        Base64.getDecoder().decode(key))
                        );
            } else {
                return keyFactory
                        .generatePrivate(
                                new PKCS8EncodedKeySpec(
                                        Base64.getDecoder().decode(key))
                        );
            }
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new SecurityException("秘钥实例化失败", e);
        }
    }

    /**
     * 获取秘钥对.
     *
     * @param bits     ⽐特位
     * @param password 密码
     */
    public static KeyPair getKey(int bits, String password) {
        try {
            val keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(bits, password == null
                    ? new SecureRandom() : new SecureRandom(password.getBytes()));
            return keyPairGenerator.generateKeyPair();
        } catch (NoSuchAlgorithmException e) {
            throw new SecurityException(e);
        }
    }
}



